package com.yc;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.nio.file.FileSystems;
import java.nio.file.PathMatcher;
import java.util.List;

public class EasyRAGExample {

    public static void main(String[] args) {
        try {

            // 1. 加载文档 - 从指定目录加载所有 PDF 文件
            String docPath = System.getProperty("user.dir") + "/demo/src/main/java/com/yc/document";
            List<Document> documents = FileSystemDocumentLoader.loadDocuments(docPath);

            // 2. 预处理并存储文档到嵌入存储（向量数据库）
            InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
            EmbeddingStoreIngestor.ingest(documents, embeddingStore);

            // 3. 配置聊天模型
            String apiKey=System.getenv("DEEPSEEK_API_KEY");
            OpenAiChatModel model = OpenAiChatModel.builder()
                    .apiKey(apiKey)
                    .modelName("deepseek-chat")
                    .baseUrl("https://api.deepseek.com")
                    .logRequests(true)
                    .logResponses(true)
                    .build();

            

            // 4. 创建 AI 助手服务
            interface Assistant {
                String chat(String userMessage);
            }

            Assistant assistant = AiServices.builder(Assistant.class)
                    .chatModel(model)
                    .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
                    .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
                    .build();

            // 5. 与 AI 助手交互
            String userQuestion = "如何使用 LangChain4j 的 Easy RAG 功能？";
            String answer = assistant.chat(userQuestion);
            System.out.println("用户问题: " + userQuestion);
            System.out.println("AI 回答: " + answer);

        } catch (Exception e) {
            e.printStackTrace();
        }
    }
//
//    // 自定义文档加载方法（包含过滤功能）
//    private static List<Document> loadDocuments(String directoryPath, PathMatcher pathMatcher) {
//        try {
//            // 使用文件系统文档加载器加载文档，并应用路径匹配器过滤
//            return dev.langchain4j.document.loader.FileSystemDocumentLoader
//                    .loadDocuments(directoryPath, pathMatcher);
//        } catch (Exception e) {
//            throw new RuntimeException("加载文档时出错", e);
//        }
//    }
}